import networkx as nx
from typing import TypeVar, Generic, Optional, Callable, Any, Iterable

NODE_TYPE = TypeVar("NODE_TYPE")
EDGE_TYPE = TypeVar("EDGE_TYPE")


class NetworkXKG(nx.DiGraph, Generic[NODE_TYPE, EDGE_TYPE]):

    def __init__(self, incoming_graph_data=None, **attr):
        super().__init__(incoming_graph_data, **attr)
        self._node_types_list: list[NODE_TYPE] = []
        self._edge_types_list: list[EDGE_TYPE] = []

        # 从节点类型、函数名映射到节点的id
        self.node_type_name_mapping: dict[NODE_TYPE, dict[str, int]] = {}

        # 不同类型的节点的id集合
        self.node_types_index: dict[NODE_TYPE, set[int]] = {}

        # 不同类型节点的标签格式化
        # {节点类型： ((节点id，节点数据) -> 标签内容)}
        self.node_label_formatters: dict[NODE_TYPE,
                                         Callable[[int, Any], str]] = {}
        # {边类型： ((边起点id，边终点id，边数据) -> 标签内容)}
        self.edge_label_formatters: dict[EDGE_TYPE,
                                         Callable[[int, int, Any], str]] = {}

    @classmethod
    def create(cls, node_types_list: list[str], edge_types_list: list[str]):
        """
        一般情况下必须要使用这个方法创建。
        """
        kg = cls()
        for node_type in node_types_list:
            kg.node_type_name_mapping[node_type] = {}
            kg.node_types_index[node_type] = set()
        kg._node_types_list = node_types_list
        kg._edge_types_list = edge_types_list
        return kg

    def _update_indices(self):
        """
        根据内部的节点和边的信息，重新计算节点/边类型的索引
        """
        all_node_types = set([self.nodes[n]['type'] for n in self.nodes])
        self.node_type_name_mapping = {k: {}
                                       for k in all_node_types}
        self.node_types_index = {k: set()
                                 for k in all_node_types}
        for node_id in self.nodes:
            node_data = self.nodes[node_id]
            node_type = node_data['type']
            node_name = node_data['name']
            self.node_type_name_mapping[node_type][node_name] = node_id
            self.node_types_index[node_type].add(node_id)

    def add_node_label_formatter(self, node_type: NODE_TYPE, node_formatter: Callable[[int, Any], str]):
        """
        添加节点标签的格式化函数
        """
        self.node_label_formatters[node_type] = node_formatter

    def add_edge_label_formatter(self, edge_type: EDGE_TYPE, edge_formatter: Callable[[int, int, Any], str]):
        """
        添加边标签的格式化函数
        """
        self.edge_label_formatters[edge_type] = edge_formatter

    def subgraph(self, nodes: Iterable):
        """
        生成子图，同时复制属性。
        """
        subgraph: NetworkXKG = super().subgraph(nodes)
        subgraph.edge_label_formatters = self.edge_label_formatters
        subgraph.node_label_formatters = self.node_label_formatters
        subgraph._update_indices()
        return subgraph

    def add_edge(self, u_of_edge, v_of_edge, edge_type: EDGE_TYPE, **attr):
        return super().add_edge(u_of_edge, v_of_edge, type=edge_type, **attr)

    def ensure_node(
        self, node_type: NODE_TYPE, name: str = "", **attr
    ) -> Optional[int]:
        mapping = self.node_type_name_mapping[node_type]
        if name not in mapping:
            node_id = self.new_node_id()
            self.add_node(node_id, type=node_type, name=name, **attr)
            self.node_types_index[node_type].add(node_id)
            mapping[name] = node_id
            return mapping[name]
        return mapping[name]

    def get_node_id(self, node_type: NODE_TYPE, name: str = "") -> int:
        return self.node_type_name_mapping[node_type].get(name)

    def new_node_id(self):
        if "node_id_max" not in self.graph:
            self.graph["node_id_max"] = 0
        self.graph["node_id_max"] += 1
        return self.graph["node_id_max"]

    def nodes_of_type(self, node_type: NODE_TYPE) -> set[int]:
        return self.node_types_index[node_type]

    def format_node_label(self, node_id: int, node_data: Any) -> str:
        node_type: NODE_TYPE = node_data['type']
        if node_type in self.node_label_formatters:
            return self.node_label_formatters[node_type](node_id, node_data)
        else:
            return node_data.get("name", str(node_id))

    def format_edge_label(self,
                          edge: tuple[int, int], edge_data: Any
                          ):
        edge_type: EDGE_TYPE = edge_data['type']
        if edge_type in self.edge_label_formatters:
            return self.edge_label_formatters[edge_type](*edge, edge_data)
        else:
            return edge_data.get("label", "")
